# -*- coding:UTF-8 -*-

from __future__ import print_function
from speech.ibm.speech_recognition import IBMSpeechRecognizer
from translate.youdao.translate_recognition import YouDaoTranslateRecognizer
import threading
import librosa
from vad.speech_segmentation import compute_bic, vad
import speech_recognition
from queue import Queue
from tools.common_tools import toTimeStr


class TextData:
    def __init__(self, start_time, end_time, en_text=None, zh_text=None):
        self.start_time = start_time
        self.end_time = end_time
        self.en_text = en_text
        self.zh_text = zh_text


class Subtitler(threading.Thread):
    def __init__(self, times_queue, results_queue):
        super(Subtitler, self).__init__()
        self.times_queue = times_queue
        self.results_queue = results_queue
        self.file_path = None

        self.__sr = IBMSpeechRecognizer()
        self.__tr = YouDaoTranslateRecognizer()
        self.__r = speech_recognition.Recognizer()
        self.__audio = None

        self.__flag = threading.Event()  # 用于暂停线程的标识
        self.__flag.set()  # 设置为True
        self.__running = threading.Event()  # 用于停止线程的标识
        self.__running.set()  # 将running设置为True

    def set_file_path(self, file_path):
        self.file_path = file_path
        with speech_recognition.AudioFile(self.file_path) as source:
            self.__audio = self.__r.record(source)

    def run(self):
        while self.__running.isSet():
            self.__flag.wait()

            data = self.times_queue.get()
            segment_audio = self.__audio.get_segment(data.start_time*1000, data.end_time*1000)
            username, password = self.__sr.username, self.__sr.password

            # 语音识别
            try:
                en_text = self.__r.recognize_ibm(segment_audio, username, password, 'en-US')
                if en_text:
                    data.en_text = en_text.replace('\n', ',')
                    # 语音翻译
                    data.zh_text = self.__tr.translate(data.en_text)
            except Exception as e:
                    print(e)
                    print('%s -> %s 未完成' % (data.start_time, data.end_time))

            self.results_queue.put(data)

    def pause(self):
        self.__flag.clear()  # 设置为False, 让线程阻塞

    def resume(self):
        self.__flag.set()  # 设置为True, 让线程停止阻塞

    def stop(self):
        self.__flag.set()  # 将线程从暂停状态恢复, 如何已经暂停的话
        self.__running.clear()  # 设置为False


class Filter(threading.Thread):
    """
        音频线程类
    """
    def __init__(self, times_queue, sr, frame_size, frame_shift):
        super(Filter, self).__init__()
        self.times_queue = times_queue
        self.sr = sr
        self.frame_size = frame_size
        self.frame_shift = frame_shift

        self.start_point = 0
        self.end_point = 200
        self.grow_point = 200
        self.delta = 25
        self.temp_point = self.start_point
        self.file_path = None
        self.y = None
        self.mfccs = None
        self.mfccs_n = None

        self.__flag = threading.Event()  # 用于暂停线程的标识
        self.__flag.set()  # 设置为True
        self.__running = threading.Event()  # 用于停止线程的标识
        self.__running.set()  # 将running设置为True

    def set_file_path(self, file_path):
        self.file_path = file_path
        # -> 时间级数和
        y, sr = librosa.load(self.file_path, sr=self.sr)

        # -> mfcc数组
        mfccs = librosa.feature.mfcc(y, self.sr, n_mfcc=12, hop_length=self.frame_shift, n_fft=self.frame_size)
        mfccs = mfccs / mfccs.max()
        m, n = mfccs.shape

        self.y = y
        self.mfccs = mfccs
        self.mfccs_n = n

    def change_start_point(self, time):
        self.start_point = int(time * float(self.sr) / self.frame_shift)
        self.end_point = self.start_point + self.grow_point
        self.temp_point = self.start_point

    def run(self):
        while self.__running.isSet():
            self.__flag.wait()

            if self.end_point < self.mfccs_n:
                feature_seg = self.mfccs[:, self.start_point:self.end_point]
                det_bic = compute_bic(feature_seg, self.delta)  # 计算BIC

                if det_bic > 0:
                    temp = self.start_point + det_bic

                    start_time, end_time = self.time_quantum(self.y, self.temp_point, temp)
                    if start_time and end_time:
                        self.times_queue.put(TextData(start_time, end_time))
                    self.temp_point = temp

                    self.start_point = self.start_point + det_bic + self.grow_point
                    self.end_point = self.start_point + self.grow_point
                else:
                    self.end_point = self.end_point + self.grow_point
            elif self.start_point < self.mfccs_n <= self.end_point and self.temp_point < self.mfccs_n:
                start_time, end_time = self.time_quantum(self.y, self.temp_point, self.mfccs_n)
                if start_time and end_time:
                    self.times_queue.put(TextData(start_time, end_time))

    def time_quantum(self, y, start_point, end_point):
        """
        获取需要处理的语音时间段
            :param y: 时间序列
            :param start_point: mfcc数组开始标记
            :param end_point: mfcc数组结束标记
            :param sr: 目标取样率
            :param frame_size:
            :param frame_shift: mfcc数组到时间序列的映射值
            :return:
        """
        start_point = start_point * self.frame_shift
        end_point = end_point * self.frame_shift
        temp_seg = y[start_point:end_point]

        start_time = end_time = None

        x1, x2 = vad.vad(temp_seg, sr=self.sr, frame_size=self.frame_size, frame_shift=self.frame_shift)
        if len(x1) != 0 and len(x2) != 0:
            start_time = start_point / float(self.sr)
            end_time = end_point / float(self.sr)

        return start_time, end_time

    def pause(self):
        self.__flag.clear()  # 设置为False, 让线程阻塞

    def resume(self):
        self.__flag.set()  # 设置为True, 让线程停止阻塞

    def stop(self):
        self.__flag.set()  # 将线程从暂停状态恢复, 如何已经暂停的话
        self.__running.clear()  # 设置为False


class Handler(threading.Thread):
    """
        操作线程方法
    """
    def __init__(self, times_queue, results_queue, filter, subtitler):
        super(Handler, self).__init__()
        self.times_queue = times_queue
        self.results_queue = results_queue
        self.filter = filter
        self.subtitler = subtitler

        self.__flag = threading.Event()  # 用于暂停线程的标识
        self.__flag.set()  # 设置为True
        self.__running = threading.Event()  # 用于停止线程的标识
        self.__running.set()  # 将running设置为True

    def run(self):
        index = 1
        while self.__running.isSet():
            self.__flag.wait()

            if index > 10:
                self.filter.pause()
                self.subtitler.pause()
                with self.times_queue.mutex:
                    self.times_queue.queue.clear()
                with self.results_queue.mutex:
                    self.results_queue.queue.clear()
                self.filter.change_start_point(0)
                self.filter.resume()
                self.subtitler.resume()
                index = 0
            else:
                data = self.results_queue.get()
                print(str(index) + "：" + toTimeStr(data.start_time) + " -> " + toTimeStr(data.end_time))
                if data.en_text and data.zh_text:
                    print(data.en_text + " -> " + data.zh_text)
                print()

            index += 1

    def pause(self):
        self.__flag.clear()  # 设置为False, 让线程阻塞

    def resume(self):
        self.__flag.set()  # 设置为True, 让线程停止阻塞

    def stop(self):
        self.__flag.set()  # 将线程从暂停状态恢复, 如何已经暂停的话
        self.__running.clear()  # 设置为False

if __name__ == '__main__':
    frame_size = 256
    frame_shift = 128
    sr = 16000
    file_path = "E:/迅雷下载/videoplayback.wav"

    q1 = Queue(50)
    q2 = Queue(50)

    t1 = Filter(q1, sr, frame_size, frame_shift)
    t2 = Subtitler(q1, q2)
    tt = Handler(q1, q2, t1, t2)

    t1.set_file_path(file_path)
    t2.set_file_path(file_path)

    t1.start()
    t2.start()
    tt.start()

    tt.join()